package com.rymcu.forest.openai.service.impl;

import com.rymcu.forest.openai.service.SseService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Created on 2023/5/26 11:38.
 *
 * @author ronger
 * @email ronger-x@outlook.com
 * @desc : com.rymcu.forest.openai.service.impl
 */
@Slf4j
@Service
public class SseServiceImpl implements SseService {



    private static final Map<Long, SseEmitter> sessionMap = new ConcurrentHashMap<>();

    @Override
    public SseEmitter connect(Long idUser) {
        if (existsUser(idUser)) {
            removeUser(idUser);
        }
        SseEmitter sseEmitter = new SseEmitter(0L);
        sseEmitter.onError((err) -> {
            log.error("type: SseSession Error, msg: {} session Id : {}", err.getMessage(), idUser);
            onError(idUser, err);
        });

        sseEmitter.onTimeout(() -> {
            log.info("type: SseSession Timeout, session Id : {}", idUser);
            removeUser(idUser);
        });

        sseEmitter.onCompletion(() -> {
            log.info("type: SseSession Completion, session Id : {}", idUser);
            removeUser(idUser);
        });
        addUser(idUser, sseEmitter);
        return sseEmitter;
    }

    @Override
    public boolean send(Long idUser, String content) {
        if (existsUser(idUser)) {
            try {
                SseEmitter emitter = sessionMap.get(idUser);
                if (emitter != null) {
                    // 检查连接是否还活跃
                    try {
                        emitter.send(SseEmitter.event().data(content).name("message"));
                        log.debug("成功向用户 {} 发送SSE消息: {}", idUser, content.substring(0, Math.min(content.length(), 50)) + "...");
                        return true;
                    } catch (Exception e) {
                        log.error("SSE连接已断开，用户 {}: {}", idUser, e.getMessage());
                        removeUser(idUser);
                        return false;
                    }
                }
            } catch (Exception exception) {
                log.error("type: SseSession send Error: {}, msg: {} session Id : {}", exception.getClass().getSimpleName(), exception.getMessage(), idUser);
                // 连接出现异常，移除用户会话
                removeUser(idUser);
            }
        } else {
            log.warn("尝试向未建立SSE连接的用户 {} 发送消息: {}", idUser, content.substring(0, Math.min(content.length(), 50)) + "...");
        }
        return false;
    }

    @Override
    public void close(Long idUser) {
        log.info("type: SseSession Close, session Id : {}", idUser);
        removeUser(idUser);
    }

    private void addUser(Long idUser, SseEmitter sseEmitter) {
        sessionMap.put(idUser, sseEmitter);
    }

    private void onError(Long sessionKey, Throwable throwable) {
        SseEmitter sseEmitter = sessionMap.get(sessionKey);
        if (sseEmitter != null) {
            sseEmitter.completeWithError(throwable);
        }
    }

    private void removeUser(Long idUser) {
        sessionMap.remove(idUser);
    }

    private boolean existsUser(Long idUser) {
        return sessionMap.containsKey(idUser);
    }

    private void sendMessage(Long idUser, String content) throws IOException {
        sessionMap.get(idUser).send(content);
    }

    /**
     * 心跳检测 - 每30秒发送一次心跳
     */
    @Scheduled(fixedRate = 30000)
    public void heartbeat() {
        if (sessionMap.isEmpty()) {
            return;
        }

        log.debug("发送SSE心跳，当前活跃连接数: {}", sessionMap.size());

        sessionMap.forEach((idUser, emitter) -> {
            try {
                emitter.send(SseEmitter.event().data("heartbeat").name("ping"));
                log.debug("向用户 {} 发送心跳成功", idUser);
            } catch (Exception e) {
                log.warn("向用户 {} 发送心跳失败，移除连接: {}", idUser, e.getMessage());
                removeUser(idUser);
            }
        });
    }

    /**
     * 获取当前活跃连接数
     */
    public int getActiveConnections() {
        return sessionMap.size();
    }
}
